package 并查集_下;

import java.util.Arrays;

public class Code03_NumberOfGoodPaths {

    public static int MAXN = 30001;

    // 需要保证集合中，代表节点的值，一定是整个集合的最大值
    public static int[] father = new int[MAXN];

    // 集合中最大值的次数，也就是 集合中代表节点的值有几个
    public static int[] maxcnt = new int[MAXN];

    public static void build(int n){
        for(int i = 0; i < n; i++){
            father[i] = i;
            maxcnt[i] = 1;
        }
    }

    // 这个并查集的优化只来自扁平化
    public static int find(int i){
        if(i != father[i]){
            father[i] = find(father[i]);
        }
        return father[i];
    }

    // 核心！
    // 注意以下的写法！
    // 谁的值大，谁做代表节点
    // 同时注意 maxcnt 的更新
    public static int union(int x, int y, int[] vals){
        // fx : x所在集团的代表节点，同时也是x所在集团的最大值下标
        int fx = find(x);
        // fy : y所在集团的代表节点，同时也是y所在集团的最大值下标
        int fy = find(y);
        // 产生好路径的条数
        int path = 0;
        if(vals[fx] > vals[fy]){ //没有好路径
            father[fy] = fx;
        } else if (vals[fx] < vals[fy]) { //没有好路径
            father[fx] = fy;
        }else {// 两个集合的最大值一样
            // 产生好路径的条数
            path = maxcnt[fx] * maxcnt[fy];
            // fx做合并后集合的代表节点
            father[fy] = fx;
            maxcnt[fx] += maxcnt[fy];
        }
        return path;
    }

    public static int numberOfGoodPaths(int[] vals, int[][] edges) {
        int n = vals.length;
        build(n);
        int ans = n;
        Arrays.sort(edges,
                (e1, e2) -> (Math.max(vals[e1[0]],
                        vals[e1[1]]) - Math.max(vals[e2[0]], vals[e2[1]])));
        for (int[] edge : edges) {
            ans += union(edge[0], edge[1], vals);
        }
        return ans;
    }

    // 课上讲解的例子1和例子2
    public static void main(String[] args) {
        // 课上例子1
        //              0  1  2  3  4  5  6  7
        int[] vals1 = { 2, 1, 1, 2, 2, 1, 1, 2 };
        int[][] edges1 = {
                { 0, 1 },
                { 0, 2 },
                { 1, 3 },
                { 2, 4 },
                { 2, 5 },
                { 5, 6 },
                { 6, 7 } };
        System.out.println(numberOfGoodPaths(vals1, edges1));

        // 课上例子2
        //              0  1  2  3  4  5  6  7  8  9 10 11 12
        int[] vals2 = { 1, 2, 2, 3, 1, 2, 2, 1, 1, 3, 3, 3, 3 };
        int[][] edges2 = {
                { 0, 1 },
                { 0, 2 },
                { 0, 3 },
                { 1, 4 },
                { 4, 7 },
                { 4, 8 },
                { 3, 5 },
                { 3, 6 },
                { 6, 9 },
                { 6, 10 },
                { 6, 11 },
                { 9, 12 } };
        System.out.println(numberOfGoodPaths(vals2, edges2));
    }
}
